use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::super::TypeEnvironment;
use super::super::TypeInference;

impl TypeInference {



    /// Infer types in a call expression with expected type
    pub(in crate::type_check::type_inference) fn infer_call_expression_with_expected(
        &mut self,
        call: &CallExpr,
        env: &mut TypeEnvironment,
        expected_type: Option<&Type>,
    ) -> Result<Type> {
        // Handle special case: Result::Ok(value) or Result::Err(error)
        if let Expression::MethodCall(method_call) = &*call.callee {
            if let Expression::Path(path_expr) = &*method_call.object {
                if path_expr.segments.len() == 2 && path_expr.segments[0] == "Result" {
                    let variant_name = &method_call.method;

                    // Handle Result::Ok(value) constructor
                    if variant_name == "Ok" {
                        if call.arguments.len() != 1 {
                            return Err(CompilerError::type_error(
                                call.location.line,
                                call.location.column,
                                format!(
                                    "Result::Ok() expects 1 argument, got {}",
                                    call.arguments.len()
                                ),
                            ));
                        }

                        // Infer the argument type - this becomes the Ok type
                        let ok_type = self.infer_expression(&call.arguments[0], env)?;

                        // Try to infer Err type from context (expected type)
                        let err_type = if let Some(Type::Result(expected_result)) = expected_type {
                            // If expected type is Result<T, E>, use E as the error type
                            *expected_result.err_type.clone()
                        } else {
                            // Otherwise, use string as the default error type
                            Type::Basic(crate::ast::types::BasicType::String)
                        };

                        return Ok(Type::Result(crate::ast::types::ResultType {
                            ok_type: Box::new(ok_type),
                            err_type: Box::new(err_type),
                            location: path_expr.location,
                        }));
                    }

                    // Handle Result::Err(error) constructor
                    if variant_name == "Err" {
                        if call.arguments.len() != 1 {
                            return Err(CompilerError::type_error(
                                call.location.line,
                                call.location.column,
                                format!(
                                    "Result::Err() expects 1 argument, got {}",
                                    call.arguments.len()
                                ),
                            ));
                        }

                        // Infer the argument type - this becomes the Err type
                        let err_type = self.infer_expression(&call.arguments[0], env)?;

                        // Try to infer Ok type from context (expected type)
                        let ok_type = if let Some(Type::Result(expected_result)) = expected_type {
                            // If expected type is Result<T, E>, use T as the ok type
                            *expected_result.ok_type.clone()
                        } else {
                            // Otherwise, use void as the default ok type
                            Type::Basic(crate::ast::types::BasicType::Void)
                        };

                        return Ok(Type::Result(crate::ast::types::ResultType {
                            ok_type: Box::new(ok_type),
                            err_type: Box::new(err_type),
                            location: path_expr.location,
                        }));
                    }
                }
            }
        }

        // Infer callee type
        let callee_type = self.infer_expression(&call.callee, env)?;
        if let Type::Basic(basic) = &callee_type {
            if *basic == crate::ast::types::BasicType::Void {
                if let Expression::Variable(name) = &*call.callee {
                    if let Some(function_type) = env.lookup_function(name) {
                        return Ok(function_type
                            .return_type
                            .as_ref()
                            .map(|t| *t.clone())
                            .unwrap_or(Type::Basic(crate::ast::types::BasicType::Void)));
                    }
                }
            }
        }

        match callee_type {
            Type::Function(func_type) => {
                // Check argument count
                if call.arguments.len() != func_type.parameter_types.len() {
                    return Err(CompilerError::type_error(
                        call.location.line,
                        call.location.column,
                        format!(
                            "Expected {} arguments, got {}",
                            func_type.parameter_types.len(),
                            call.arguments.len()
                        ),
                    ));
                }

                // Infer argument types
                for (arg, param_type) in call.arguments.iter().zip(func_type.parameter_types.iter())
                {
                    let arg_type = self.infer_expression(arg, env)?;
                    self.add_constraint(&arg_type, &param_type, call.location)?;
                }

                // Return function return type
                Ok(func_type
                    .return_type
                    .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)))
                    .as_ref()
                    .clone())
            }
            _ => Err(CompilerError::type_error(
                call.location.line,
                call.location.column,
                "Cannot call non-function type",
            )),
        }
    }



    /// Infer types in a method call expression
    pub(in crate::type_check::type_inference) fn infer_method_call_expression(
        &mut self,
        method_call: &MethodCallExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        // Infer object type
        let object_type = self.infer_expression(&method_call.object, env)?;

        // Handle Result type method calls
        if let Type::Result(result_type) = &object_type {
            return self.infer_result_method_call(
                result_type,
                &method_call.method,
                &method_call.arguments,
                &method_call.location,
                env,
            );
        }

        // Handle Option type method calls
        if let Type::Option(option_type) = &object_type {
            return self.infer_option_method_call(
                option_type,
                &method_call.method,
                &method_call.arguments,
                &method_call.location,
                env,
            );
        }

        if let Type::Rc(rc_type) = &object_type {
            if method_call.method == "downgrade" {
                if !method_call.arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        method_call.location.line,
                        method_call.location.column,
                        "Rc::downgrade() does not take any arguments",
                    ));
                }

                return Ok(Type::Weak(crate::ast::types::WeakType {
                    inner_type: rc_type.inner_type.clone(),
                    location: method_call.location,
                }));
            }
        }

        if let Type::Weak(weak_type) = &object_type {
            if method_call.method == "lock" {
                if !method_call.arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        method_call.location.line,
                        method_call.location.column,
                        "Weak::lock() does not take any arguments",
                    ));
                }

                return Ok(Type::Option(crate::ast::types::OptionType {
                    some_type: Box::new(Type::Rc(crate::ast::types::RcType {
                        inner_type: weak_type.inner_type.clone(),
                        location: weak_type.location,
                    })),
                    location: method_call.location,
                }));
            }
        }

        // Try to find method in impl blocks (format: "TypeName.methodName")
        let method_name = match &object_type {
            Type::Struct(struct_type) => {
                // First try to find method with TypeName.methodName format
                format!("{}.{}", struct_type.name, method_call.method)
            }
            _ => {
                // For non-struct types, try simple method name first
                method_call.method.clone()
            }
        };

        // Look up method in function environment
        let function_type = if let Some(ft) = env.lookup_function(&method_name) {
            Some(ft)
        } else if let Type::Struct(_) = &object_type {
            // If not found with TypeName.methodName, try simple method name
            env.lookup_function(&method_call.method)
        } else {
            None
        };

        if let Some(function_type) = function_type {
            // Check argument count
            // Note: impl methods have self as first parameter, so we need to account for that
            let expected_args = if function_type.parameter_types.len() > 0 {
                function_type.parameter_types.len() - 1 // Subtract self parameter
            } else {
                function_type.parameter_types.len()
            };

            if method_call.arguments.len() != expected_args {
                return Err(CompilerError::type_error(
                    method_call.location.line,
                    method_call.location.column,
                    format!(
                        "Expected {} arguments, got {}",
                        expected_args,
                        method_call.arguments.len()
                    ),
                ));
            }

            // Clone the function type to avoid borrowing issues
            let param_types = function_type.parameter_types.clone();
            let return_type = function_type.return_type.clone();

            // Infer argument types
            // Skip the first parameter (self) if it exists
            // Note: impl methods have self as first parameter, so we need to skip it
            let param_types_to_check = if param_types.len() > method_call.arguments.len() {
                // First param is self, skip it
                &param_types[1..]
            } else {
                &param_types[..]
            };

            for (arg, param_type) in method_call
                .arguments
                .iter()
                .zip(param_types_to_check.iter())
            {
                let arg_type = self.infer_expression(arg, env)?;
                self.add_constraint(&arg_type, param_type, method_call.location)?;
            }

            // Return function return type
            Ok(return_type
                .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)))
                .as_ref()
                .clone())
        } else {
            // If method not found, try to infer from static method call
            // For static method calls like Person::new, the object is the type name
            if let Expression::Variable(type_name) = &*method_call.object {
                // Look for static method in type environment
                if let Some(_type_def) = env.lookup_type(type_name) {
                    // For now, assume static methods return the type they belong to
                    // This is a simplified implementation
                    if method_call.method == "new" {
                        // Constructor returns the type it belongs to
                        Ok(Type::Struct(StructType {
                            name: type_name.clone(),
                            type_args: vec![],
                            location: method_call.location,
                        }))
                    } else {
                        // Other static methods return void for now
                        Ok(Type::Basic(crate::ast::types::BasicType::Void))
                    }
                } else {
                    Err(CompilerError::type_error(
                        method_call.location.line,
                        method_call.location.column,
                        format!("Unknown type: {}", type_name),
                    ))
                }
            } else {
                Err(CompilerError::type_error(
                    method_call.location.line,
                    method_call.location.column,
                    format!(
                        "Unknown method: {} for type {:?}",
                        method_call.method, object_type
                    ),
                ))
            }
        }
    }

}
